package cgroups

import (
	"bufio"
	"os"
	"path"
	"strings"
)

const mountinfo = "/proc/self/mountinfo"

type ResourceLimit struct {
	Memory string
	Cpu    string
	CpuSet string
}

type CgroupManager struct {
	SubPath string
	ResourceLimit
}

var (
	SubSystems = []SubSystem{
		&MemorySubSystem{},
	}
)

type SubSystem interface {
	// Name 返回cgroup子系统名
	Name() string

	// Set 设置资源限制
	Set(subPath string, res *ResourceLimit) error

	// Apply 添加进程到cgroups
	Apply(subPath string, pid int) error

	// Remove 移除
	Remove(subPath string) error
}

func NewCgroupManager(path string) *CgroupManager {
	return &CgroupManager{
		SubPath: path,
	}
}

func (c *CgroupManager) Apply(pid int) error {
	for _, subSystem := range SubSystems {
		subSystem.Apply(c.SubPath, pid)
	}
	return nil
}

func (c *CgroupManager) Set(limit *ResourceLimit) error {
	for _, subSystem := range SubSystems {
		subSystem.Set(c.SubPath, limit)
	}
	return nil
}

func (c *CgroupManager) Destroy() error {
	for _, subSystem := range SubSystems {
		subSystem.Remove(c.SubPath)
	}
	return nil
}

func GetGroupPath(SubPath string, subSystemName string, autoCreate bool) (wholeCgroupPath string, err error) {
	cgroupMountPoint := findCgroupMountPoint(subSystemName) // /sys/fs/cgroup/memory

	wholeCgroupPath = path.Join(cgroupMountPoint, SubPath)

	if _, err = os.Stat(wholeCgroupPath); err == nil {
		return
	}

	if autoCreate && os.IsNotExist(err) {
		err = os.Mkdir(wholeCgroupPath, 0755)
	}
	return
}

func findCgroupMountPoint(subSystemName string) string {
	file, err := os.Open(mountinfo)
	if err != nil {
		return ""
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		text := scanner.Text()
		split := strings.Split(text, " ")
		for _, name := range strings.Split(split[len(split)-1], ",") {
			if name == subSystemName {
				return split[4]
			}
		}

		if err = scanner.Err(); err != nil {
			return ""
		}
	}
	return ""
}
